/*
 * Copyright (c) 2025 Industrial Software Feature Database
 */
package com.comac.ins.test;

import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import org.bson.BsonBinaryReader;
import org.bson.BsonDocument;
import org.bson.ByteBufNIO;
import org.bson.Document;
import org.bson.codecs.BsonDocumentCodec;
import org.bson.codecs.DecoderContext;
import org.bson.io.ByteBufferBsonInput;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;

@SpringBootTest
public class BsonImportTest {

    @Autowired
    private MongoClient mongoClient;

    @Test
    void importLargeBsonFile() throws IOException {
        // === 1. 文件 & collection 名称 ===
        String filePath = "C:\\Users\\lijpt\\Desktop\\op2_aircraft.bson";
        File bsonFile = new File(filePath);
        String collectionName = bsonFile.getName().replaceFirst("\\.bson$", "");

        // === 2. 如果DB不存在，自动创建 ===
        MongoDatabase database = mongoClient.getDatabase("simright");
        MongoCollection<org.bson.Document> collection = database.getCollection(collectionName);

        // 如果 collection 已存在，可以先清理（可选）
        collection.drop();
        collection = database.getCollection(collectionName);

        // === 3. 批量写入参数 ===
        int batchSize = 10000; // 每批 1w 条
        List<Document> buffer = new ArrayList<>(batchSize);

        // === 4. 逐条读取 BSON ===
        try (FileInputStream fis = new FileInputStream(bsonFile)) {
            byte[] lengthBytes = new byte[4];
            BsonDocumentCodec codec = new BsonDocumentCodec();

            while (true) {
                // 先读文档长度
                int read = fis.read(lengthBytes);
                if (read < 0) break;

                int docLength = ((lengthBytes[0] & 0xFF)) |
                    ((lengthBytes[1] & 0xFF) << 8) |
                    ((lengthBytes[2] & 0xFF) << 16) |
                    ((lengthBytes[3] & 0xFF) << 24);

                byte[] fullDoc = new byte[docLength];
                System.arraycopy(lengthBytes, 0, fullDoc, 0, 4);
                int remaining = docLength - 4;
                int actuallyRead = fis.readNBytes(fullDoc, 4, remaining);
                if (actuallyRead < remaining) break; // EOF

                // 解析 BSON -> BsonDocument
                try (BsonBinaryReader reader = new BsonBinaryReader(
                    new ByteBufferBsonInput(new ByteBufNIO(ByteBuffer.wrap(fullDoc))))) {
                    BsonDocument bsonDoc = codec.decode(reader, DecoderContext.builder().build());
                    buffer.add(org.bson.Document.parse(bsonDoc.toJson()));
                }

                // 到达 batchSize 就写入一次
                if (buffer.size() >= batchSize) {
                    collection.insertMany(buffer);
                    buffer.clear();
                }
            }

            // 插入最后剩余的
            if (!buffer.isEmpty()) {
                collection.insertMany(buffer);
            }
        }

        System.out.println("导入完成: " + collectionName + " 共写入数据。");
    }
}
